
#include "xdns-handler.h"
#include "xdns-address-query.h"
#include "xdns-data.h"
#include "xdns-decoder.h"
#include "xdns-header-info.h"
#include "xdns-host-entry-generator.h"
#include "xdns-name-query.h"
#include "xdns-protocol.h"
#include "xdns-question-info.h"
#include "xdns-resource-record-info.h"
#include "xdns-server-state.h"
#include <arpa/inet.h>
#include <assert.h>
#include <string.h>
#include <limits>
#include <sstream>

using namespace std;

const long XDNS_MAX_QUERY_ID = numeric_limits<long>::max();

XdnsHandler::XdnsHandler()
{
}

XdnsHandler::~XdnsHandler()
{
}

XdnsError
XdnsHandler::Init(const XdnsConfig& config,
                  XdnsHandler** handler,
                  string* errorDescription)
{
    if (NULL == handler)
    {
        if (errorDescription != NULL)
        {
            (*errorDescription) = "Parameter handler is NULL.";
        }

        return XE_Inval;
    }

    // Convert the Internet host address from a string.
    struct in_addr serverAddress = {0};
    serverAddress.s_addr = inet_addr(config.m_serverIpAddress.c_str());
    if (INADDR_NONE == serverAddress.s_addr)
    {
        if (errorDescription != NULL)
        {
            (*errorDescription) = "Server IP address ";
            (*errorDescription) += config.m_serverIpAddress;
            (*errorDescription) += " is invalid.";
        }

        return XE_Inval;
    }

    XdnsServerState* serverState = NULL;
    XdnsError error = XdnsServerState::Init(config.m_serverPort,
                                            &serverAddress,
                                            &serverState);
    if (error != XE_Ok)
    {
        if (errorDescription != NULL)
        {
           (*errorDescription) = "Failed to initialize the state of a DNS server.";
        }

        return error;
    }

    XdnsHandlerImpl* handlerImpl = new (nothrow) XdnsHandlerImpl(serverState);
    if (NULL == handlerImpl)
    {
        XdnsServerState::Release(serverState);
        serverState = NULL;

        if (errorDescription != NULL)
        {
           (*errorDescription) = "No enough memory could be used to create a XDNS handler.";
        }

        return XE_NoMem;
    }

    *handler = handlerImpl;
    return XE_Ok;
}

void
XdnsHandler::Release(XdnsHandler* handler)
{
    if (handler != NULL)
    {
        delete handler;
    }
}

XdnsHandlerImpl::XdnsHandlerImpl(XdnsServerState* serverState)
: XdnsHandler(),
  m_nextQueryId(XDNS_MAX_QUERY_ID),
  m_serverState(serverState),
  m_lastError(),
  m_queries()
{
    assert(m_serverState != NULL);
    m_serverState->RegisterListener(*this);
}

XdnsHandlerImpl::~XdnsHandlerImpl()
{
    if (m_serverState != NULL)
    {
        m_serverState->UnregisterListener(*this);
        XdnsServerState::Release(m_serverState);
    }

    map<long, XdnsQuery*>::iterator iter = m_queries.begin();
    map<long, XdnsQuery*>::iterator end = m_queries.end();
    for (; iter != end; ++iter)
    {
        delete (iter->second);
    }
    m_queries.clear();
}

XdnsError
XdnsHandlerImpl::GetHostByName(const string& name,
                               XdnsListener* listener,
                               long* queryId)
{
    long newQueryId = 0;
    map<long, XdnsQuery*>::iterator newIter = m_queries.end();
    XdnsError error = GetValidQueryId(newQueryId, newIter);
    if (error != XE_Ok)
    {
        m_lastError = "Cannot retrieve a valid query ID, error = ";
        m_lastError += LongToString(static_cast<long>(error));
        m_lastError += ".";
        return error;
    }

    XdnsNameQuery* query = new (nothrow) XdnsNameQuery(name, listener);
    if (NULL == query)
    {
        m_lastError = "No enough memory could be used to create a name query.";
        return XE_NoMem;
    }

    assert(m_serverState != NULL);
    error = m_serverState->SendQuery(newQueryId, *query);
    if (error != XE_Ok)
    {
        delete query;
        m_lastError = "An error occurred when sending a query, error = ";
        m_lastError += LongToString(static_cast<long>(error));
        m_lastError += ".";
        return error;
    }

    m_queries.insert(newIter, make_pair(newQueryId, query));
    if (queryId != NULL)
    {
        *queryId = newQueryId;
    }
    return XE_Ok;
}

XdnsError
XdnsHandlerImpl::GetHostByAddress(const struct in_addr* ipAddress,
                                  XdnsListener* listener,
                                  long* queryId)
{
    if (NULL == ipAddress)
    {
        return XE_Inval;
    }

    long newQueryId = 0;
    map<long, XdnsQuery*>::iterator newIter = m_queries.end();
    XdnsError error = GetValidQueryId(newQueryId, newIter);
    if (error != XE_Ok)
    {
        m_lastError = "Cannot retrieve a valid query ID, error = ";
        m_lastError += LongToString(static_cast<long>(error));
        m_lastError += ".";
        return error;
    }

    XdnsAddressQuery* query = new (nothrow) XdnsAddressQuery(ipAddress, listener);
    if (NULL == query)
    {
        m_lastError = "No enough memory could be used to create an address query.";
        return XE_NoMem;
    }

    assert(m_serverState != NULL);
    error = m_serverState->SendQuery(newQueryId, *query);
    if (error != XE_Ok)
    {
        delete query;
        m_lastError = "An error occurred when sending a query, error = ";
        m_lastError += LongToString(static_cast<long>(error));
        m_lastError += ".";
        return error;
    }

    m_queries.insert(newIter, make_pair(newQueryId, query));
    if (queryId != NULL)
    {
        *queryId = newQueryId;
    }
    return XE_Ok;
}

void
XdnsHandlerImpl::CancelQuery(long queryId)
{
    map<long, XdnsQuery*>::iterator iter = m_queries.find(queryId);
    if (iter != m_queries.end())
    {
        delete (iter->second);
        m_queries.erase(iter);
    }

    assert(m_serverState != NULL);
    m_serverState->CancelQuery(queryId);
}

XdnsError
XdnsHandlerImpl::GetProcessingFds(int* nfds,
                                  fd_set* readfds,
                                  fd_set* writefds)
{
    assert(m_serverState != NULL);
    XdnsError error = m_serverState->GetProcessingFds(nfds, readfds, writefds);
    if ((error != XE_Ok) &&
        (error != XE_NoEnt))
    {
        m_lastError = "An error occurred when retrieving the processing set of file descriptors, error = ";
        m_lastError += LongToString(static_cast<long>(error));
        m_lastError += ".";
        return error;
    }

    return error;
}

XdnsError
XdnsHandlerImpl::Process(fd_set* readfds,
                         fd_set* writefds)
{
    assert(m_serverState != NULL);
    XdnsError error = m_serverState->Process(readfds, writefds);
    if (error != XE_Ok)
    {
        m_lastError = "An error occurred when processing the activated file descriptors, error = ";
        m_lastError += LongToString(static_cast<long>(error));
        m_lastError += ".";
        return error;
    }

    return XE_Ok;
}

string
XdnsHandlerImpl::GetLastErrorDescription() const
{
    return m_lastError;
}

void
XdnsHandlerImpl::ReplyReceived(long queryId,
                               XdnsError error,
                               const XdnsData* data,
                               size_t startOffset,
                               const XdnsHeaderInfo* header)
{
    map<long, XdnsQuery*>::iterator queryIter = m_queries.find(queryId);
    if (queryIter == m_queries.end())
    {
        // The query is canceled.
        return;
    }
    XdnsQuery* query = queryIter->second;
    assert(query != NULL);

    // Retrieve the requested hostname or IP address.
    string requestedHostname;
    const struct in_addr* requestedIpAddress = NULL;
    XdnsQueryType queryType = query->GetType();
    switch (queryType)
    {
    case XQT_Hostname:
    {
        XdnsNameQuery* nameQuery = static_cast<XdnsNameQuery*>(query);
        requestedHostname = nameQuery->GetHostname();
        assert(!(requestedHostname.empty()));
        break;
    }
    case XQT_IpAddress:
    {
        XdnsAddressQuery* addressQuery = static_cast<XdnsAddressQuery*>(query);
        requestedIpAddress = addressQuery->GetAddress();
        assert(requestedIpAddress != NULL);
        break;
    }
    default:
    {
        assert(false);
        break;
    }
    }

    vector<XdnsHostEntry*> hostEntries;
    do
    {
        if (error != XE_Ok)
        {
            break;
        }

        // Check if the data is correct.
        if (NULL == data)
        {
            error = XE_Unexpected;
            break;
        }

        // Decode the header info if it is NULL.
        size_t offset = startOffset;
        XdnsHeaderInfo tempHeader;
        if (NULL == header)
        {
            error = XdnsDecoder::DecodeHeader(*data, offset, tempHeader);
            if (error != XE_Ok)
            {
                break;
            }
            header = &tempHeader;
        }
        assert(header != NULL);

        int answerCount = static_cast<int>(header->m_ancount);
        if ((header->m_qdcount != 1) || (answerCount <= 0))
        {
            error = XE_BadResponse;
            break;
        }

        // Decode the question info.
        offset += XDNS_HEADER_SIZE;
        XdnsQuestionInfo question;
        size_t questionSize = 0;
        error = XdnsDecoder::DecodeQuestion(*data,
                                            startOffset,
                                            offset,
                                            question,
                                            questionSize);
        if (error != XE_Ok)
        {
            break;
        }

        /* Check if the question info of response is same with the question
           info retrieved from the query.
        */
        XdnsQuestionInfo requestedQuestion;
        query->GetQuestion(requestedQuestion);
        if ((question.m_qtype != requestedQuestion.m_qtype) ||
            (question.m_qclass != requestedQuestion.m_qclass) ||
            (question.m_qname.compare(requestedQuestion.m_qname) != 0))
        {
            error = XE_BadResponse;
            break;
        }

        // Decode the answers of the resource record info.
        offset += questionSize;
        hostEntries.reserve(static_cast<size_t>(answerCount));
        while ((--answerCount) >= 0)
        {
            XdnsResourceRecordInfo resourceRecord;
            size_t resourceRecordSize = 0;
            error = XdnsDecoder::DecodeResourceRecord(*data,
                                                      startOffset,
                                                      offset,
                                                      resourceRecord,
                                                      resourceRecordSize);
            if (error != XE_Ok)
            {
                break;
            }

            // Generate the host entry.
            XdnsHostEntry* hostEntry = NULL;
            error = XdnsHostEntryGenerator::Generate(requestedIpAddress,
                                                     requestedHostname,
                                                     question,
                                                     resourceRecord,
                                                     *data,
                                                     startOffset,
                                                     &hostEntry);
            if (error != XE_Ok)
            {
                /* TODO: Some formats of DNS resource record are not
                         related with the host entry and not supported.
                */
                error = XE_Ok;
                continue;
            }

            assert(hostEntry != NULL);
            hostEntries.push_back(hostEntry);

            offset += resourceRecordSize;
        }
    }
    while (0);

    if ((XE_Ok == error) && (hostEntries.empty()))
    {
        error = XE_NoEnt;
    }

    query->NotifyListener(queryId,
                          error,
                          (XE_Ok == error) ? (&hostEntries) : NULL);

    // Destroy the host entries.
    vector<XdnsHostEntry*>::iterator hostEntryIter = hostEntries.begin();
    vector<XdnsHostEntry*>::iterator hostEntryEnd = hostEntries.end();
    for (; hostEntryIter != hostEntryEnd; ++hostEntryIter)
    {
        delete (*hostEntryIter);
    }
}

/*! Retrieve a new valid query ID.

    @param queryId On return a valid query ID.
    @param iter On return an iterator used to insert the query.
    @return XE_Ok if success,
            XE_NoEnt if the query ID is collided with IDs in map.
*/
XdnsError
XdnsHandlerImpl::GetValidQueryId(long& queryId,
                                 map<long, XdnsQuery*>::iterator& iter)
{
    long newQueryId = NextQueryId();
    map<long, XdnsQuery*>::iterator newIter = m_queries.lower_bound(newQueryId);
    if ((newIter != m_queries.end()) &&
        (newIter->first == newQueryId))
    {
        /* The new query ID is present, reset the next query ID and retrieve
           another one.
        */
        m_nextQueryId = XDNS_MAX_QUERY_ID;
        newQueryId = NextQueryId();
        newIter = m_queries.lower_bound(newQueryId);
        if ((newIter != m_queries.end()) &&
           (newIter->first == newQueryId))
        {
            return XE_NoEnt;
        }
    }

    queryId = newQueryId;
    iter = newIter;
    return XE_Ok;
}

long
XdnsHandlerImpl::NextQueryId()
{
    return (m_nextQueryId--);
}

string
XdnsHandlerImpl::LongToString(long value)
{
    stringstream ss;
    ss << value;
    return ss.str();
}

